package shop.gateway.filter;

import com.fasterxml.jackson.databind.ObjectMapper;
import lombok.extern.slf4j.Slf4j;
import org.reactivestreams.Publisher;
import org.springframework.cloud.gateway.filter.GatewayFilterChain;
import org.springframework.cloud.gateway.filter.GlobalFilter;
import org.springframework.core.io.buffer.DataBuffer;
import org.springframework.core.io.buffer.DataBufferFactory;
import org.springframework.core.io.buffer.DataBufferUtils;
import org.springframework.http.server.reactive.ServerHttpResponseDecorator;
import org.springframework.stereotype.Component;
import org.springframework.web.server.ServerWebExchange;
import reactor.core.publisher.Flux;
import reactor.core.publisher.Mono;
import shop.common.core.Result;

import java.io.IOException;

@Component
@Slf4j
public class ResponseRewriteFilter implements GlobalFilter {

    private final ObjectMapper objectMapper;

    public ResponseRewriteFilter(ObjectMapper objectMapper) {
        this.objectMapper = objectMapper;
    }

    @Override
    public Mono<Void> filter(ServerWebExchange exchange, GatewayFilterChain chain) {
        // 保存原始响应体
        ServerHttpResponseDecorator decoratedResponse = new ServerHttpResponseDecorator(exchange.getResponse()) {
            @Override
            public Mono<Void> writeWith(Publisher<? extends DataBuffer> body) {
                if (body instanceof Flux) {
                    Flux<? extends DataBuffer> fluxBody = (Flux<? extends DataBuffer>) body;
                    return super.writeWith(fluxBody.buffer().map(dataBuffers -> {
                        // 合并所有数据缓冲区
                        DataBufferFactory bufferFactory = exchange.getResponse().bufferFactory();
                        DataBuffer joined = bufferFactory.join(dataBuffers);
                        byte[] content = new byte[joined.readableByteCount()];
                        joined.read(content);
                        DataBufferUtils.release(joined);

                        try {
                            // 解析原始响应
                            Result originalResult = objectMapper.readValue(content, Result.class);

                            // 如果code不是成功状态，记录日志
                            if (originalResult.getCode() != 200) {
                                log.error("下游服务返回错误: {}", originalResult);
                            }

                            // 可以在这里统一修改响应内容
                            Result modifiedResult = originalResult; // 或根据业务需要修改

                            // 返回修改后的响应
                            byte[] newContent = objectMapper.writeValueAsBytes(modifiedResult);
                            return bufferFactory.wrap(newContent);
                        } catch (IOException e) {
                            log.error("响应处理失败", e);
                            return bufferFactory.wrap(content); // 返回原始内容
                        }
                    }));
                }
                return super.writeWith(body);
            }
        };

        return chain.filter(exchange.mutate().response(decoratedResponse).build());
    }
}